#include <cstring>
#include <iostream>
#include <cmath>
#include <cstdio>


#include "Node.h"
#include "TetrahedronState.h"
#include "Tetrahedron.h"
#include "BigVector.h"
#include "FEMModel.h"
#include "Vector.h"
#include "TopologicalOperation.h"
#include "NonlinearCGState.h"
#include "DeformationState.h"
#include "Statistics.h"
#include "Misc.h"


// #define PRINT_BETA_CHOICE
// #define PRINT_NEWTON

bool use_hestenes_stiefel;
bool use_polak_ribiere;
bool use_fletcher_reeves;
Real newton_tolerance = 1e-2;

NonlinearCGState::NonlinearCGState ()
{
	no_more_improvement_b_ = false;
	restart_b_ = true;
	residual_len_ = prev_residual_len_ =0;
	iter_count_ = 0;
	last_restart_ =0;
	residual_tolerance_ = GetNumberSetting ("outer-loop-tolerance");

	char const * line_search = GetStringSetting ("nonlinear-cg-linesearch");
	if (!strcmp (line_search, "secant"))
		step_length_function_ = &NonlinearCGState::SecantMethod;
	else if (!strcmp (line_search, "linear"))
		step_length_function_ = &NonlinearCGState::LinearAlpha;
	else
		step_length_function_ = &NonlinearCGState::NewtonMethod;

	char const * beta_func = GetStringSetting ("nonlinear-cg-betafunc");
	if (!strcmp(beta_func, "polak-ribiere"))
	{
		printf ("Using PR beta selection.");    
		use_polak_ribiere= true;
		use_fletcher_reeves = false;
	}
	if (!strcmp(beta_func, "hestenes-stiefel"))
	{
		printf ("Using HS beta selection.");    
		use_hestenes_stiefel = true;
		use_polak_ribiere= false;
		use_fletcher_reeves = false;
	}
	else if (!strcmp (beta_func, "fletcher-reeves"))
	{
		printf ("Using FR beta selection.");
		use_fletcher_reeves = true;
		use_polak_ribiere = false;
	}
	else
	{
		printf ("Using both PR and FR.");
		use_polak_ribiere = use_fletcher_reeves = true;
	}

	newton_tolerance = GetNumberSetting ("nonlinear-cg-newton-tolerance");
}



/*
COST:

8*n flops
*/
void NonlinearCGState::IterationStep (DeformationState  *pstate) 
{
	int n = pstate->Dimension ();

	flop_count +=  8* n;

	global_iteration_count ++;
	Real beta =  0.0;


	/*
	In the past we had some troubles with Polak-Ribiere where
	Fletcher-Reeves performed better. They turned out to be related to
	programming errors -- optimizing an `incorrect' function gives
	interesting results.

	*/
	bool polak_ribiere = false;


	if (!restart_b_)
	{
		if (use_polak_ribiere) {
			/*
			Polak-Ribiere method.
			*/
			beta = (residual_len_  -
				BigVectorIP (residual_.AccessArray (),
				prev_residual_.AccessArray (), n)) / prev_residual_len_;
			polak_ribiere = true;
		}

		if (use_hestenes_stiefel)
		{
			beta=
				(residual_len_ - BigVectorIP (residual_.AccessArray (), prev_residual_.AccessArray (),n) )/
				(BigVectorIP(search_dir_.AccessArray (), prev_residual_.AccessArray(), n)

				- BigVectorIP (search_dir_.AccessArray(), residual_.AccessArray () , n ));

		}

		if (use_fletcher_reeves)
		{
			/*
			Fletcher-Reeves method -- can never result in negative beta.
			*/
			if (beta <= 0.0)
			{
				beta = residual_len_ / prev_residual_len_;
				polak_ribiere = false;
			}
		}

		if (beta < 0)
		{
			beta = 0.0; // restarting
			cout << "backing up!";
		}
	}
	bool restarted = ! beta;
	if (restarted)
	{
		printf (" *** restarting. Residual force = %lf\n", residual_len_);
		last_restart_ = iter_count_;
	}

	restart_b_ = false;
	iter_count_ ++;

	Real *dir = search_dir_.UnsafeAccessArray ();
	Real *disp =  pstate->displacement_.UnsafeAccessArray ();  
	bool const *fixed = pstate->fixed_arr_.AccessArray();
	Real *residu =residual_.UnsafeAccessArray ();

	BigVectorSaxpy (dir, beta, dir, residu,n);
	BigVectorPartialNullify (dir, dir, fixed, n);

	Real alpha;



#ifdef PRINT_BETA_CHOICE 
	if (polak_ribiere)
	{
		printf ("Polak-Ribiere\n");
	}
	else
	{
		printf ("Fletcher-Reeves\n");
	}
#endif

	alpha = (*step_length_function_)(this, pstate);

	if (alpha)
	{
		BigVectorSaxpy (disp, alpha, dir, disp,n );
	}
	else if (restarted)
	{
		no_more_improvement_b_ = true;
	}
	else
	{
		printf ("would have restarted. \n");
		//       restart_b_ = true;
	}
}


int NonlinearCGState::Dimension ()const
{
	return try_residual_.Size ();
}

/*
Calculate residual and 2nd derivative for SEARCH_DIR, leaving results in

TRY_DISPLACEMENT
MAT_TIMES_SEARCH_DIR
TRY_RESIDUAL_

and  lengths in

DERIV, RESLEN
*/
void
NonlinearCGState::CalculateResidualAndDerivative (Real *reslen,
													   Real * deriv,
													   Real alpha, DeformationState const*pstate)
{
	Real *matdir = mat_times_search_dir_.UnsafeAccessArray ();
	Real *try_disp =  try_displacement_.UnsafeAccessArray ();
	Real * try_residual = try_residual_.UnsafeAccessArray ();
	Real * dir = search_dir_.UnsafeAccessArray ();
	Real const * disp = pstate->displacement_.AccessArray ();
	int n =Dimension ();
	BigVectorSaxpy  (try_disp, alpha, dir, disp, n );
	BigVectorNullify (matdir, n);
	BigVectorNullify (try_residual, n);
	pstate->ApplyForceDerivativeAndResidual (matdir, try_residual, try_disp, dir);

	*reslen = BigVectorIP (try_residual, try_residual, n);

	/*
	residual_derivative = d ||residual||/d alpha
	=  / ( residual * DForce[dir] ) / ||residual||
	*/
	*deriv = BigVectorIP (try_residual, matdir, n)  / sqrt(*reslen);
}


/*
Interestingly, nonlinear CG doesn't make the residual length go down
monotonously. We should look at the energy measure instead, -- that
should go down.

*/

Real NonlinearCGState::NewtonMethod (NonlinearCGState *me,
								   DeformationState const *pstate)
{
	int j_max  = 4;
	int j = 0;
	int n =me->Dimension ();
	Real alpha = 0.0;
	Real const *d = me->search_dir_.AccessArray ();
	Real delta_0 = 0.0;
	Real dirlen = sqrt (BigVectorIP (d,d,n));
	while (j < j_max)
	{
		Real res, der;

		me->CalculateResidualAndDerivative (&res, &der, alpha, pstate);

		Real delta = - BigVectorIP (me->try_residual_.AccessArray (), d, n)
			/ BigVectorIP (d, me->mat_times_search_dir_.AccessArray (), n);
		if (j == 0)
			delta_0 = delta;


		alpha +=delta;
#ifdef PRINT_NEWTON
		cout << "iter " << j << " alpha = " << alpha << " delta =  " << delta << " dirlen " << dirlen << endl;
#endif

		if (delta  < newton_tolerance * delta_0)
			break;

		/*
		Ugh. This is a non-scalable stopping criterion. I should be ashamed!

		Anyway, we normally converge in 1 step. Theoretically, the
		result of the last CalculateResidualAndDerivative () could
		be reused for the update of the real residual force.

		Perhaps the elastic force should be computed and stored
		separately. It only changes when the displacements are
		changed.

		*/
		if (delta * dirlen <  newton_tolerance)
			break ; 
		j++;
	}

	return alpha;
}

Real NonlinearCGState::CalculateTryResidual (Real alpha, DeformationState const *pstate)
{
	Real *try_disp =  try_displacement_.UnsafeAccessArray ();
	Real * try_residual = try_residual_.UnsafeAccessArray ();
	Real * dir = search_dir_.UnsafeAccessArray ();
	Real const * disp = pstate->displacement_.AccessArray ();

	int n =Dimension ();
	BigVectorSaxpy  (try_disp, alpha, dir, disp, n );
	BigVectorNullify (try_residual, n);  
	return pstate->ComputeResidualForce (try_residual, try_disp);
}

Real NonlinearCGState::LinearAlpha (NonlinearCGState *me, DeformationState const*pstate)
{
	Real * matdir = me->mat_times_search_dir_.UnsafeAccessArray();
	Real const *dir  = me->search_dir_.AccessArray();

	pstate->ApplyForceDerivative (matdir, 0,
		pstate->displacement_.AccessArray (),
		dir);

	int n =me->Dimension();
	Real alpha = - BigVectorIP (me->residual_.AccessArray (), dir, n) / BigVectorIP (matdir, dir, me->Dimension( ));

	return alpha;
}



const Real SECANT_TOLERANCE = 1e-6;

Real NonlinearCGState::SecantMethod (NonlinearCGState *me,
								   DeformationState const*pstate)
{
	int n = me->Dimension();

	Real *dir = me->search_dir_.UnsafeAccessArray ();
	Real *res =  me->try_residual_.UnsafeAccessArray ();

	Real last_alpha = -0.0;
	//  CalculateTryResidual (last_alpha, pstate);

	Real last_eta = BigVectorIP (me->residual_.AccessArray (),  dir, n);
	Real alpha = 1.0;

	//  me->graph_alpha (1.0, pstate);


	Real dirlen = BigVectorIP (dir, dir, n);
	Real delta = 1.0;
	int j_max = 5;
	int j = 0;
	while (j++ < j_max && dirlen * fabs (delta)  > Sqr (SECANT_TOLERANCE))
	{
		//           Real reslen =
		me->CalculateTryResidual (alpha, pstate);      
		Real eta =  BigVectorIP (res, dir,n);

		/*
		cout << "eta, leta " << eta << " " << last_eta << endl;      
		cout << "alpha,lalpha " << alpha << " "  << last_alpha<<endl;
		*/

		delta = (eta * (last_alpha - alpha)) / (last_eta - eta);

		last_alpha = alpha;
		alpha = alpha - delta;

		//cout << "new alpha " <<alpha<<endl;
		last_eta = eta;  
	}


	return alpha;
}



void NonlinearCGState::UpdateResidual (DeformationState const*pstate)
{
	bool *fixp  = (bool*) pstate->fixed_arr_.AccessArray ();
	Real *delta_f = (Real*) pstate->delta_external_force_.AccessArray ();

	prev_residual_len_ =residual_len_;
	int n = pstate->Dimension ();

	prev_residual_  = residual_;
	residual_len_ = pstate->ComputeResidualForce (residual_.UnsafeAccessArray (),
		pstate->displacement_.AccessArray ());

	bool change = false;
	for (int i = n; i--;)
	{
		if (*delta_f && !*fixp)
		{
			change = true;
		}

		delta_f ++;
		fixp ++;
	}


	if (isnan (residual_len_) || isinf (residual_len_))
	{
		cerr << "Overflow. Exit";
		abort ();
	}

	int restart_count = (int) GetNumberSetting ("nonlinear-cg-restart-count");
	if (restart_count < 0)
		restart_count = n;

	if (change
		|| (restart_count && (iter_count_ -last_restart_ >  restart_count))
		)
	{
		restart_b_ = true;
		no_more_improvement_b_ = false;
		/*
		we don't set iter_count_ = 0, since it messes up the
		convergence stats.

		*/
	}
}



/*
Compute the residual, put result into RESIDUAL_VECTOR. Return the
square of the length.  */

/*
It seems nonsensical to keep looking if RESIDUAL_LEN_ is too small,
unfortunately, it makes measuring convergence impossible, since it
introduces a scaling parameter.  */
bool NonlinearCGState::GoodSolution (DeformationState const * pstate)const
{
	if (no_more_improvement_b_)
		return true;

	bool b = residual_len_ <= pstate->external_force_len_sq_ * Sqr (residual_tolerance_);

	if (!(pstate->external_force_len_sq_))
	{
		b = b || residual_len_ < Sqr (residual_tolerance_);
	}

	return b;
}

void NonlinearCGState::PrintNodes () const
{
}

void NonlinearCGState::Print () const
{
	cout << "res "<<residual_len_ <<endl;
}

/*
Synchronize all state and precomputed information, with the current static
mesh. 
*/
void NonlinearCGState::Reinit (DeformationState const * pstate) 
{
	int n =Dimension ();
	BigVectorNullify (search_dir_.UnsafeAccessArray(), n);
	BigVectorNullify (mat_times_search_dir_.UnsafeAccessArray(), n);

	/*
	Updated from update_forces
	*/
	residual_len_ = 0.0;

	no_more_improvement_b_ = false;
	prev_residual_len_ =0.0;
	iter_count_ =0;
	restart_b_ = true;
}


extern bool Contains (Array<int> const * is, int j);

LinkArray< Array<Real> > NonlinearCGState::VectorVariables ()
{
	Array<Real> *vecs[]
	= {
		&mat_times_search_dir_, &search_dir_, &residual_, &try_displacement_, &try_residual_, &prev_residual_, 0,
	};

	LinkArray< Array<Real> > vvs;
	for (int i= 0; vecs[i]; i++)
		vvs.Push (vecs[i]);
	return vvs;
}

NonlinearCGState::~NonlinearCGState ()
{
}


/****************************************************************/



void NonlinearCGState::GraphAlpha (Real start_step, DeformationState const*pstate)
{
	cout << "start step "  << start_step << endl;
	int n = pstate->Dimension();
	static int graphcount ;
	char filename[1024];

	printf("graph %d\n", graphcount);
	sprintf (filename, "tab%d", graphcount);

	FILE *gr =fopen (filename, "w");
	// plot graph.
	const int STEP_LIMIT= 40;
	Real start_graph= -0.5*start_step;
	Real stop_graph = 1.1*start_step;
	Real alpha = start_graph;
	Real delta = ((stop_graph - start_graph) / STEP_LIMIT); 
	int j =0;
	Real last_res = 0;
	do
	{
		Real res, deriv;
		CalculateResidualAndDerivative (&res, &deriv, alpha, pstate);

		Real new_residual = pstate->ComputeResidualForce (try_residual_.UnsafeAccessArray (),
			try_displacement_.UnsafeAccessArray ());
		flop_count += 2 *n ;

		if (j)
		{
			Real deriv_est =  (sqrt (new_residual) - sqrt (last_res)) / delta;
			fprintf (gr, "%f %f\n", alpha,  deriv_est);
		}
		last_res = new_residual;
		alpha += delta;
	}
	while (j ++ < STEP_LIMIT);
	fclose (gr);

	char s[1024];
	sprintf (s, "gnuplot %s", filename);
	system (s);

	graphcount++;
}
